import numpy as np
from assignment2.python.helper import camera2


# def chooseP2(E, Pts2):
#     # load in the essential matrix and get 4 possible P2 camera matrix
#     M2s = camera2(E)
#
#     add = np.ones((1, Pts2.shape[0]), dtype=np.float32, order='C')
#     pts2 = np.insert(Pts2, 2, values=add, axis=1)
#
#     # prepare for counting
#     count = np.zeros(4, dtype=int, order='C')
#
#     k_load = np.load('assignment2/data/intrinsics.npz')
#     k1, k2 = k_load['K1'], k_load['K2']
#
#     # print(M2s)
#     for i in range(0, 4):
#         M = M2s[:, :, i]
#         #print(M)
#         R = M[0:3, 0:3].reshape(3, 3)
#         #print(np.linalg.det(R))
#         if np.abs(np.linalg.det(R) - 1) < 0.001:
#             print("good one")
#             # get real_3D points N*4 matrix
#             # real_pts = (np.linalg.inv(k1 @ M) @ pts2.T).T
#             # for dots in real_pts:
#             #     if dots[2] < 0:
#             count[i] += 1
#
#     print(count)
#
#     P2 = M2s[:, :, np.argmin(count)].reshape(3, 4)
#     # print(P2)
#     # print(P2)
#     return P2

def get_camera_matrices(p2, k1, k2):
    p1 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0]], dtype=np.float32)
    P1 = k1 @ p1
    P2 = k2 @ p2
    return P1, P2


def help_mat_create(pt1, pt2, P1, P2):
    # print(pt1[0] * P1[2] - P1[0])
    A = np.array([
        pt1[0] * P1[2] - P1[0],
        pt1[1] * P1[2] - P1[1],
        pt2[0] * P2[2] - P2[0],
        pt2[1] * P2[2] - P2[1]
    ], dtype=np.float32)
    # print(A)
    return A


def triangulate(P1, pts1, P2, pts2):
    pts3 = np.empty((pts1.shape[0], 3), dtype=float, order='C')
    # print(pts1.shape)
    # print(pts2.shape)

    for i in range(0, pts1.shape[0]):
        A = help_mat_create(pts1[i], pts2[i], P1, P2)
        # print(A.shape)
        u, s, v = np.linalg.svd(A)
        P_dash = v[-1]
        P = [P_dash[0] / P_dash[3], P_dash[1] / P_dash[3], P_dash[2] / P_dash[3]]
        # print(P)
        pts3[i] = P

    return pts3


def project_back(Pts3, p1, pts1):
    add = np.ones((1, Pts3.shape[0]), dtype=np.float32, order='C')
    pts3 = np.insert(Pts3, 3, values=add, axis=1)
    # print(pts3)
    pts3_2d = (p1 @ pts3.T).T
    distance = 0
    # print(pts3_2d)
    for i in range(0, pts3_2d.shape[0]):
        distance += np.sqrt(
            (pts3_2d[i][0] / pts3_2d[i][2] - pts1[i][0]) ** 2 + (pts3_2d[i][1] / pts3_2d[i][2] - pts1[i][1]) ** 2)
    print(distance)
